import torch.nn as nn

from .sste import SparseSuperTokenEncoder_OpenPCDet
from .img_backbones import Swin_OpenPCDet, SWIN_CFG



def build_img_model(img_cfg):
    """
    a image model takes batch_dict as input and outputs to image_features key
    """
    # reference REF_NECK_CFG in swin.py to get neck cfg
    return Swin_OpenPCDet(SWIN_CFG)


def build_lidar_model(lidar_cfg):
    """
    a lidar model takes batch_dict as input and outputs to pillar_features and voxel_coords key
    """
    return SparseSuperTokenEncoder_OpenPCDet(lidar_cfg)


class MixedTR(nn.Module):
    """
    mixing different transformer backbones
    """
    def __init__(self, model_cfg, **kwargs):
        super().__init__()
        
        img_cfg = model_cfg.image_cfg
        lidar_cfg = model_cfg.lidar_cfg
        
        self.img_model = build_img_model(img_cfg)
        self.lidar_model = build_lidar_model(lidar_cfg)
        
    def forward(self, batch_dict):
        batch_dict = self.img_model(batch_dict)
        batch_dict = self.lidar_model(batch_dict)
        return batch_dict